import heapq
from typing import Dict, Optional

class HuffmanNode:
    """哈夫曼树的节点"""
    def __init__(self, char: Optional[str], freq: int):
        self.char = char 
        self.freq = freq
        self.left = None
        self.right = None

    def __lt__(self, other):
        """用于优先队列比较，频率小的优先"""
        return self.freq < other.freq

class HuffmanTree:
    """哈夫曼树类"""
    def __init__(self):
        self.root = None
        self.codes = {}

    def __str__(self) -> str:
        """返回哈夫曼树的编码字典字符串"""
        if not self.codes:
            self.get_codes()
        return str(self.codes)
    
    def _generate_codes(self, node: HuffmanNode, current_code: str) -> None:
        """递归生成编码字典"""
        if not node:
            return

        if node.char is not None:
            self.codes[node.char] = current_code
            return

        self._generate_codes(node.left, current_code + '0')
        self._generate_codes(node.right, current_code + '1')

    def build_tree(self, freq_dict: Dict[str, int]) -> None:
        """根据字频字典构建哈夫曼树

        Args:
            freq_dict (Dict[str, int]): 字频字典，其中键为字符，值为其出现的频率。

        Returns:
            None: 直接更新类的 `self.root` 属性，表示构建的哈夫曼树根节点。
        """
        heap = [HuffmanNode(char, freq) for char, freq in freq_dict.items()]
        heapq.heapify(heap)

        while len(heap) > 1:
            node1 = heapq.heappop(heap)
            node2 = heapq.heappop(heap)

            merged = HuffmanNode(None, node1.freq + node2.freq)
            merged.left = node1
            merged.right = node2

            heapq.heappush(heap, merged)

        self.root = heap[0] if heap else None

    def get_codes(self) -> Dict[str, str]:
        """根据哈夫曼树生成编码字典"""
        self.codes = {}
        self._generate_codes(self.root, '')
        return self.codes

if __name__ == "__main__":
    freq_dict = {'a': 5, 'b': 9, 'c': 12, 'd': 13, 'e': 16, 'f': 45}

    huffman_tree = HuffmanTree()
    huffman_tree.build_tree(freq_dict)
    print(huffman_tree)